package com.bungarus.codec;

import com.bungarus.model.SessionBuilder;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.StringTokenizer;

/**
 * Created by tang on 2019/2/12.
 */
public class TokenGenerator {
    static int RSA_KEY_LENGTH = 2048;
    static String ALGORITHM_NAME = "RSA" ;
    static String PADDING_SCHEME = "OAEPWITHSHA-512ANDMGF1PADDING" ;
    static String MODE_OF_OPERATION = "ECB" ; // This essentially means none behind the scene

    public static void main(String args[]) {
        String shortMessage = args[0] ;

        try {
            // Generate Key Pairs
            KeyPairGenerator rsaKeyGen = KeyPairGenerator.getInstance(ALGORITHM_NAME) ;
            rsaKeyGen.initialize(RSA_KEY_LENGTH) ;
            KeyPair rsaKeyPair = rsaKeyGen.generateKeyPair() ;
            String publicKeyStr = Base64.getEncoder().encodeToString(rsaKeyPair.getPublic().getEncoded());
            String privateKeyStr = Base64.getEncoder().encodeToString(rsaKeyPair.getPrivate().getEncoded());
            String aesKey = SessionBuilder.KeyGenerator.generate(16);
            String aesIV = SessionBuilder.KeyGenerator.generate(16);
            System.out.println("Public Key: " + publicKeyStr);
            System.out.println("Private Key: " + privateKeyStr);

            String midEncryptedText = aesEncrypt(shortMessage, aesKey, aesIV);
            String midEnRsaText = rsaEncrypt(aesKey+"."+aesIV, rsaKeyPair.getPrivate());
            String encryptedText = midEncryptedText + "." + midEnRsaText;

            String decryptKeyIV = rsaDecrypt(Base64.getDecoder().decode(midEnRsaText), rsaKeyPair.getPublic());
            StringTokenizer tokenizer = new StringTokenizer(decryptKeyIV, ".");
            aesKey = tokenizer.nextToken();
            aesIV = tokenizer.nextToken();

            String decryptedText = aesDecrypt(Base64.getDecoder().decode(midEncryptedText), aesKey, aesIV);

            System.out.println("Encrypted text = " + encryptedText) ;
            System.out.println("Decrypted text = " + decryptedText) ;

        } catch(Exception e) {System.out.println("Exception while encryption/decryption") ;e.printStackTrace() ; }


    }

    public static String generate(String plainText, Key privateKey) throws Exception {
        String aesKey = SessionBuilder.KeyGenerator.generate(16);
        String aesIV = SessionBuilder.KeyGenerator.generate(16);
        String encryptedMessagePart1 = aesEncrypt(plainText, aesKey, aesIV);
        String encryptedMessagePart2 = rsaEncrypt(aesKey+"."+aesIV, privateKey);
        return encryptedMessagePart1 + "." + encryptedMessagePart2;
    }

    public static String rsaEncrypt(String message, Key privateKey) throws Exception {

        Cipher c = Cipher.getInstance(ALGORITHM_NAME);//Cipher.getInstance(ALGORITHM_NAME + "/" + MODE_OF_OPERATION + "/" + PADDING_SCHEME) ;

        c.init(Cipher.ENCRYPT_MODE, privateKey) ;

        byte[] cipherTextArray = c.doFinal(message.getBytes("utf-8")) ;

        return Base64.getEncoder().encodeToString(cipherTextArray) ;

    }


    public static String rsaDecrypt(byte[] encryptedMessage, Key publicKey) throws Exception {
        Cipher c = Cipher.getInstance(ALGORITHM_NAME);//Cipher.getInstance(ALGORITHM_NAME + "/" + MODE_OF_OPERATION + "/" + PADDING_SCHEME) ;
        c.init(Cipher.DECRYPT_MODE, publicKey);
        byte[] plainText = c.doFinal(encryptedMessage);

        return new String(plainText, "utf-8") ;

    }

    public static String aesEncrypt(String message, String key, String iv) throws Exception {
        SecretKeySpec skeySpec = new SecretKeySpec(key.getBytes("utf-8"), "AES");
        Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");

        IvParameterSpec ips = new IvParameterSpec(iv.getBytes("utf-8"));
        cipher.init(Cipher.ENCRYPT_MODE, skeySpec, ips);
        byte[] encrypted = cipher.doFinal(message.getBytes("utf-8"));
        return Base64.getEncoder().encodeToString(encrypted);
    }

    public static String aesDecrypt(byte[] encryptedMessage, String key, String iv) throws Exception {
        try {
            byte[] raw = key.getBytes("utf-8");
            SecretKeySpec skeySpec = new SecretKeySpec(raw, "AES");
            Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
            IvParameterSpec ips = new IvParameterSpec(iv.getBytes("utf-8"));
            cipher.init(Cipher.DECRYPT_MODE, skeySpec, ips);
            try {
                byte[] original = cipher.doFinal(encryptedMessage);
                String originalString = new String(original, "utf-8");
                return originalString;
            } catch (Exception e) {
                System.out.println(e.toString());
                return null;
            }
        } catch (Exception ex) {
            System.out.println(ex.toString());
            return null;
        }
    }
}
